{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8e7cf0c9",
   "metadata": {},
   "source": [
    "# 文本分类任务\n",
    "\n",
    "本notebook用于完成文本分类任务，主要包含以下步骤：\n",
    "1. 数据加载和预处理\n",
    "2. 特征提取\n",
    "3. 模型训练\n",
    "4. 模型评估\n",
    "5. 预测和结果保存"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98adf706",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.feature_extraction.text import TfidfVectorizer\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.metrics import f1_score, precision_score, recall_score\n",
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModel\n",
    "from tqdm import tqdm\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aebb1a4a",
   "metadata": {},
   "source": [
    "## 1. 数据加载和预处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91ad3721",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_jsonl(file_path):\n",
    "    data = []\n",
    "    with open(file_path, 'r', encoding='utf-8') as f:\n",
    "        for line in f:\n",
    "            data.append(json.loads(line))\n",
    "    return data\n",
    "\n",
    "# 加载训练数据\n",
    "train_data = load_jsonl('train.jsonl')\n",
    "print(f\"训练集大小: {len(train_data)}\")\n",
    "\n",
    "# 转换为DataFrame\n",
    "train_df = pd.DataFrame(train_data)\n",
    "print(\"\\n数据示例:\")\n",
    "print(train_df.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f88ab2d",
   "metadata": {},
   "source": [
    "## 2. 特征提取"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f9c5a7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TextFeatureExtractor:\n",
    "    def __init__(self):\n",
    "        self.tfidf = TfidfVectorizer(max_features=5000)\n",
    "        # 使用本地bert-base-chinese模型路径\n",
    "        local_model_path = \"./bert-base-chinese\"\n",
    "        self.tokenizer = AutoTokenizer.from_pretrained(local_model_path)\n",
    "        self.model = AutoModel.from_pretrained(local_model_path)\n",
    "        self.model.to('cuda')  # 将模型移动到GPU\n",
    "        self.model.eval()      # 设置为推理模式，节省显存\n",
    "    \n",
    "    def extract_tfidf_features(self, texts):\n",
    "        return self.tfidf.fit_transform(texts)\n",
    "    \n",
    "    def extract_bert_features(self, texts, batch_size=64, max_length=512):\n",
    "        features = []\n",
    "        texts = list(texts)  # 保证输入为List[str]\n",
    "        for i in tqdm(range(0, len(texts), batch_size)):\n",
    "            batch_texts = texts[i:i + batch_size]\n",
    "            inputs = self.tokenizer(batch_texts, padding=True, truncation=True, max_length=max_length, return_tensors=\"pt\")\n",
    "            inputs = {k: v.to('cuda') for k, v in inputs.items()}  # 输入也放到GPU\n",
    "            with torch.no_grad():\n",
    "                outputs = self.model(**inputs)\n",
    "            # 支持多卡自动释放显存\n",
    "            features.append(outputs.last_hidden_state.mean(dim=1).cpu().numpy())  # 回到CPU\n",
    "            del inputs, outputs\n",
    "            torch.cuda.empty_cache()\n",
    "        return np.vstack(features)\n",
    "\n",
    "# 特征提取\n",
    "extractor = TextFeatureExtractor()\n",
    "X_tfidf = extractor.extract_tfidf_features(train_df['text'])\n",
    "X_bert = extractor.extract_bert_features(list(train_df['text']), batch_size=64, max_length=512)\n",
    "y = train_df['label']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff9eacf2",
   "metadata": {},
   "source": [
    "## 3. 模型训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9339d974",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TextClassifier:\n",
    "    def __init__(self):\n",
    "        self.tfidf_model = LogisticRegression(max_iter=1000)\n",
    "        self.bert_model = LogisticRegression(max_iter=1000)\n",
    "        \n",
    "    def train(self, X_tfidf, X_bert, y):\n",
    "        # 训练TF-IDF模型\n",
    "        self.tfidf_model.fit(X_tfidf, y)\n",
    "        # 训练BERT模型\n",
    "        self.bert_model.fit(X_bert, y)\n",
    "    \n",
    "    def predict(self, X_tfidf, X_bert):\n",
    "        # 集成两个模型的预测结果\n",
    "        tfidf_pred = self.tfidf_model.predict_proba(X_tfidf)\n",
    "        bert_pred = self.bert_model.predict_proba(X_bert)\n",
    "        # 加权平均\n",
    "        ensemble_pred = 0.4 * tfidf_pred + 0.6 * bert_pred\n",
    "        return (ensemble_pred[:, 1] > 0.5).astype(int)\n",
    "\n",
    "# 划分训练集和验证集\n",
    "X_tfidf_train, X_tfidf_val, X_bert_train, X_bert_val, y_train, y_val = train_test_split(\n",
    "    X_tfidf, X_bert, y, test_size=0.2, random_state=42\n",
    ")\n",
    "\n",
    "# 训练模型\n",
    "classifier = TextClassifier()\n",
    "classifier.train(X_tfidf_train, X_bert_train, y_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbae5bd8",
   "metadata": {},
   "source": [
    "## 4. 模型评估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e5b4278",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_model(y_true, y_pred):\n",
    "    precision = precision_score(y_true, y_pred)\n",
    "    recall = recall_score(y_true, y_pred)\n",
    "    f1 = f1_score(y_true, y_pred)\n",
    "    print(\"模型评估结果：\")\n",
    "    print(f\"  精确率（Precision）：{precision:.4f} —— 预测为正的样本中有多少是真的正样本\")\n",
    "    print(f\"  召回率（Recall）：{recall:.4f} —— 所有正样本中有多少被正确预测出来\")\n",
    "    print(f\"  F1分数（F1 Score）：{f1:.4f} —— 精确率和召回率的调和平均，更全面反映模型表现\")\n",
    "    if f1 > 0.8:\n",
    "        print(\"模型表现优秀！可以放心用于实际预测。\")\n",
    "    elif f1 > 0.6:\n",
    "        print(\"模型表现一般，还可以进一步优化。\")\n",
    "    else:\n",
    "        print(\"模型表现较差，建议调整特征或模型参数。\")\n",
    "\n",
    "# 在验证集评估\n",
    "val_pred = classifier.predict(X_tfidf_val, X_bert_val)\n",
    "print(\"验证集评估结果:\")\n",
    "evaluate_model(y_val, val_pred)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2d7a57b",
   "metadata": {},
   "source": [
    "## 5. 预测和结果保存"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39b7fdc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 加载测试数据\n",
    "test_data = load_jsonl('test.jsonl')\n",
    "test_df = pd.DataFrame(test_data)\n",
    "\n",
    "# 提取测试集特征\n",
    "X_test_tfidf = extractor.extract_tfidf_features(test_df['text'])\n",
    "X_test_bert = extractor.extract_bert_features(test_df['text'])\n",
    "\n",
    "# 预测\n",
    "test_pred = classifier.predict(X_test_tfidf, X_test_bert)\n",
    "\n",
    "# 保存结果\n",
    "results = pd.DataFrame({\n",
    "    'id': range(len(test_pred)),\n",
    "    'label': test_pred\n",
    "})\n",
    "results.to_csv('submission.csv', index=False)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
